import numpy as np
from mindquantum import *

#定义随机Pauli旋转门
def RP(rotation_angle):
    a = np.random.randint(0, 3)
    if a == 0:
        return RX(rotation_angle)
    elif a == 1:
        return RY(rotation_angle)
    elif a == 2:
        return RZ(rotation_angle)
    else:
        print("error in random Pauli gates")
#定义barren plateau的ansatz，线路图参考readme或者main.ipynb文件
def bpansatz(n,p):
    init_state_circ = UN(RY(np.pi/4),n)
    qc = Circuit() + init_state_circ
    for j in range(p):
        for i in range(n):
            qc += RP(f'θ({j},{i})').on(i)
        for ii in range(n-1):
            qc += Z(ii,ii+1)
        qc += Z(0, n-1)
    return qc

if __name__ == '__main__':
    bpansatz(5,1)
    print('test over')
